package cz.semjob.learning.classifier.benchmark;

import java.util.HashMap;
import java.util.Map;

public class DataSetStatistics {

	private int dataSetSize;

	private int learnDataSetSize;

	private Map<String, Integer> testAnnotationsCount;
	private Map<String, Integer> learnAnnotationsCount;

	public DataSetStatistics(int dataSetSize, int learnDataSetSize) {
		this.dataSetSize = dataSetSize;
		this.learnDataSetSize = learnDataSetSize;
		testAnnotationsCount = new HashMap<String, Integer>();
		learnAnnotationsCount = new HashMap<String, Integer>();
	}

	public void addTestAnnotationsCount(String annotation, Integer count) {
		if (testAnnotationsCount.containsKey(annotation)) {
			testAnnotationsCount.put(annotation,
					testAnnotationsCount.get(annotation) + count);
		} else {
			testAnnotationsCount.put(annotation, count);
		}
	}

	public void addLearnAnnotationsCount(String annotation, Integer count) {
		if (learnAnnotationsCount.containsKey(annotation)) {
			learnAnnotationsCount.put(annotation,
					learnAnnotationsCount.get(annotation) + count);
		} else {
			learnAnnotationsCount.put(annotation, count);
		}
	}

	public int getTotalAnnotations() {
		int total = 0;
		for (String annotation : testAnnotationsCount.keySet()) {
			total += testAnnotationsCount.get(annotation);

		}
		for (String annotation : learnAnnotationsCount.keySet()) {
			total += learnAnnotationsCount.get(annotation);

		}
		return total;
	}

	public int getDataSetSize() {
		return dataSetSize;
	}

	public int getLearnDataSetSize() {
		return learnDataSetSize;
	}

	public Map<String, Integer> getTestAnnotationsCount() {
		return testAnnotationsCount;
	}

	public Map<String, Integer> getLearnAnnotationsCount() {
		return learnAnnotationsCount;
	}

	@Override
	public String toString() {
		return "DataSetStatistics [dataSetSize=" + dataSetSize
				+ ", learnDataSetSize=" + learnDataSetSize
				+ ", testAnnotationsCount=" + testAnnotationsCount
				+ ", learnAnnotationsCount=" + learnAnnotationsCount + ", "
				+ " totalAnnotations = " + getTotalAnnotations() + " ]";
	}

}
